import os
import subprocess
from pathlib import Path
from tqdm import tqdm

def extract_audio(video_folder, output_folder, audio_format="mp3"):
    """
    Batch extract audio from multiple video files in a folder
    
    Parameters:
    video_folder (str): Path to the folder containing video files
    output_folder (str): Path to save the extracted audio files
    audio_format (str): Output audio format (default: mp3)
    """
    # Create output folder if it doesn't exist
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    
    # Get all video files in the folder and subfolders (recursive)
    video_extensions = ['.mp4', '.avi', '.mkv', '.mov', '.flv', '.wmv']
    video_files = []
    
    # Walk through all directories and subdirectories
    for root, dirs, files in os.walk(video_folder):
        for file in files:
            file_path = os.path.join(root, file)
            file_ext = os.path.splitext(file)[1].lower()
            if file_ext in video_extensions:
                video_files.append(file_path)
    
    # Extract audio from each video with progress bar
    skipped = 0
    converted = 0
    failed = 0
    
    for video_file in tqdm(video_files, desc="Extracting audio", unit="file"):
        # Preserve the subfolder structure in output
        rel_path = os.path.relpath(video_file, video_folder)
        rel_dir = os.path.dirname(rel_path)
        
        # Create subfolder in output directory if needed
        output_subdir = os.path.join(output_folder, rel_dir)
        if rel_dir and not os.path.exists(output_subdir):
            os.makedirs(output_subdir)
            
        video_filename = os.path.basename(video_file)
        audio_filename = os.path.splitext(video_filename)[0] + f".{audio_format}"
        audio_path = os.path.join(output_folder, rel_dir, audio_filename)
        
        # Skip if the audio file already exists
        if os.path.exists(audio_path):
            print(f"Skipping {video_filename} (output already exists)")
            skipped += 1
            continue
            
        # Use FFmpeg to extract audio
        command = [
            "ffmpeg",
            "-i", video_file,
            "-q:a", "0",
            "-map", "a",
            "-v", "quiet",  # Reduce FFmpeg's output verbosity
            audio_path
        ]
        
        try:
            subprocess.run(command, check=True, stderr=subprocess.PIPE)
            print(f"Successfully extracted audio to: {audio_path}")
            converted += 1
        except subprocess.CalledProcessError as e:
            print(f"Error extracting audio from {video_filename}: {str(e)}")
            failed += 1
    
    # Print summary
    print(f"\nConversion summary:")
    print(f"  Total videos: {len(video_files)}")
    print(f"  Converted: {converted}")
    print(f"  Skipped (already exist): {skipped}")
    print(f"  Failed: {failed}")

if __name__ == "__main__":
    # Example usage
    video_folder = "撩妹社群课程"  # Folder containing your videos
    output_folder = "extracted_audio"  # Folder to save extracted audio
    audio_format = "mp3"  # Change to desired format (mp3, wav, aac, etc.)
    
    extract_audio(video_folder, output_folder, audio_format)